# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Radical basis functions (RBF)"""

from typing import Union, Optional, List

import torch
from torch import nn, Tensor
import numpy as np

from ..utils import get_integer, get_tensor, get_arguments
from ..units import Length, get_length
from .rbf import RadicalBasisFunctions, _rbf_register


@_rbf_register('gaussian')
class GaussianBasis(RadicalBasisFunctions):
    """
    Gaussian radial basis functions.

    Args:
        r_max (float, optional): Maximum distance for RBF. Defaults to 1.0.
        r_min (float, optional): Minimum distance for RBF. Defaults to 0.0.
        sigma (float, optional): Standard deviation of Gaussian functions. Defaults to 0.03.
        delta (float, optional): Interval between Gaussian centers. Defaults to 0.016.
        num_basis (int, optional): Number of basis functions. Defaults to None.
        clip_distance (bool, optional): Whether to clip distance to [r_min, r_max]. Defaults to False.
        length_unit (str, optional): Unit of length. Defaults to 'nm'.
    """
    def __init__(self, 
                 r_max: float = 1.0, 
                 r_min: float = 0.0, 
                 sigma: float = 0.03, 
                 delta: Optional[float] = 0.016, 
                 num_basis: Optional[int] = None, 
                 clip_distance: bool = False, 
                 length_unit: str = 'nm',
                 **kwargs
                 ) -> None:
        super().__init__(
            r_max=r_max,
            r_min=r_min,
            num_basis=num_basis,
            clip_distance=clip_distance,
            length_unit=length_unit,
        )
        self._kwargs = get_arguments(locals(), kwargs)
        
        if num_basis is None and delta is None:
            raise TypeError('"num_basis" and "delta" cannot both be "None".')
        if num_basis is not None and num_basis <= 0:
            raise ValueError('"num_basis" must be larger than 0.')
        if delta is not None and delta <= 0:
            raise ValueError('"delta" must be larger than 0.')

        self.delta = get_tensor(delta, dtype=torch.float32, device=self.device) if delta is not None else None
        self.sigma = get_tensor(sigma, dtype=torch.float32, device=self.device)
        self.coeff = -0.5 / (self.sigma ** 2)

        if self.delta is None:
            num_intervals = int(torch.ceil(get_tensor(self.r_range / (self.r_range / (self.num_basis - 1)), dtype=torch.float32, device=self.device)) + 1)
            self.num_basis = num_intervals
            self.delta = self.r_range / (num_intervals - 1)
        else:
            if self.num_basis is None:
                num_intervals = int(torch.ceil(get_tensor(self.r_range / delta, dtype=torch.float32, device=self.device)) + 1)
                self.num_basis = num_intervals
            self.offsets = get_tensor([self.r_min + i * delta for i in range(self.num_basis)], dtype=torch.float32, device=self.device)

    def _print_info(self,
                    num_retraction: int = 6,
                    num_gap: int = 3,
                    char: str = ' '
                    ) -> None:
        ret = char * num_retraction
        gap = char * num_gap
        print(ret + gap + f' Minimum distance: {self.r_min} {self.length_unit}')
        print(ret + gap + f' Maximum distance: {self.r_max} {self.length_unit}')
        print(ret + gap + f' Interval for log Gaussian: {self.delta}')
        print(ret + gap + f' Sigma for log gaussian: {self.sigma}')
        print(ret + gap + f' Number of basis functions: {self.num_basis}')
        if self.clip_distance:
            print(ret + gap + f' Clip the range of distance to ({self.r_min}, {self.r_max}).')

    def compute_rbf(self, distance: Tensor) -> Tensor:
        if self.clip_distance:
            distance = torch.clamp(distance, self.r_min, self.r_max)

        # (..., N) -> (..., 1, N)
        ex_dis = distance.unsqueeze(-1)
        # (..., K) = (..., 1) - (K,)
        diff = ex_dis - self.offsets
        # (..., K)
        rbf = torch.exp(self.coeff * diff ** 2)

        return rbf

    def __str__(self):
        return 'GaussianBasis<>'
